package gl.drawer

import gl.util.Resource
import gl.util.nextVector3
import gl.util.putVector
import gl.vbm.VAOResult
import gl.vbm.VBMObject
import org.joml.Matrix4f
import org.joml.Vector3f
import org.joml.Vector4f
import org.lwjgl.opengl.GL45C.*
import java.nio.ByteBuffer
import kotlin.math.min
import kotlin.random.Random

/**
 * 使用 TransformFeedbacks 实现的一个粒子系统
 *
 * todo 无法运行
 */
class TransformFeedbacks : Drawer() {

    private val pointCount = 5000
    private val random = Random(0x13371337)

    private var programRender = 0
    private var uniformRndModelMatrix = 0
    private var uniformRndProjectionMatrix = 0
    private var vertexRender = 0
    private var bufferGeometry = 0
    private var texture = 0
    private lateinit var vbm: VBMObject
    private lateinit var vbmVertex: VAOResult

    private var programUpdate = 0
    private var uniformUpdModelMatrix = 0
    private var uniformUpdProjectionMatrix = 0
    private var uniformUpdTriangleCount = 0
    private var uniformUpdTimeStep = 0
    private var vertex0 = 0
    private var vertex1 = 0
    private var tbo0 = 0
    private var tbo1 = 0

    private var feedbacks = 0

    private var aspect = 1f

    override fun initialize() {
        val vertexUpdateShader = glCreateShader(GL_VERTEX_SHADER)
        glShaderSource(vertexUpdateShader, TransformFeedbacksUpdateVertex)
        glCompileShader(vertexUpdateShader)
        val fragmentUpdateShader = glCreateShader(GL_FRAGMENT_SHADER)
        glShaderSource(fragmentUpdateShader, TransformFeedbacksUpdateFragment)
        glCompileShader(fragmentUpdateShader)
        programUpdate = glCreateProgram()
        glAttachShader(programUpdate, vertexUpdateShader)
        glAttachShader(programUpdate, fragmentUpdateShader)
        glTransformFeedbackVaryings(programUpdate, arrayOf("position_out", "velocity_out"), GL_INTERLEAVED_ATTRIBS)
        glLinkProgram(programUpdate)
        glDetachShader(programUpdate, vertexUpdateShader)
        glDetachShader(programUpdate, fragmentUpdateShader)
        glDeleteShader(vertexUpdateShader)
        glDeleteShader(fragmentUpdateShader)
        val vertexRenderShader = glCreateShader(GL_VERTEX_SHADER)
        glShaderSource(vertexRenderShader, TransformFeedbacksRenderVertex)
        glCompileShader(vertexRenderShader)
        val fragmentRenderShader = glCreateShader(GL_FRAGMENT_SHADER)
        glShaderSource(vertexRenderShader, TransformFeedbacksRenderFragment)
        glCompileShader(fragmentRenderShader)
        programUpdate = glCreateProgram()
        glAttachShader(programUpdate, vertexRenderShader)
        glAttachShader(programUpdate, fragmentRenderShader)
        glTransformFeedbackVaryings(programUpdate, arrayOf("world_space_position"), GL_INTERLEAVED_ATTRIBS)
        glLinkProgram(programUpdate)
        glDetachShader(programUpdate, vertexRenderShader)
        glDetachShader(programUpdate, fragmentRenderShader)
        glDeleteShader(vertexRenderShader)
        glDeleteShader(fragmentRenderShader)

        tbo0 = glGenBuffers()
        glBindBuffer(GL_TRANSFORM_FEEDBACK_BUFFER, tbo0)
        glBufferData(GL_TRANSFORM_FEEDBACK_BUFFER, (pointCount * 28/*7*/).toLong(), GL_DYNAMIC_COPY)

        val buffer = glMapBuffer(GL_TRANSFORM_FEEDBACK_BUFFER, GL_READ_ONLY)!!
        for (i in 0 until pointCount) {
            val velocity = random.nextVector3()
            val position = Vector4f(velocity, 1.0f).add(-0.5f, 40f, 0f, 0f)
            velocity.mul(1f, 0.3f, 0.3f)
            buffer.putVector(position)
            buffer.putVector(velocity)
        }
        glUnmapBuffer(GL_TRANSFORM_FEEDBACK_BUFFER)

        vertex0 = glGenVertexArrays()
        glBindVertexArray(vertex0)
        glEnableVertexAttribArray(0)
        glVertexAttribPointer(0, 4, GL_FLOAT, false, 0, 0L)
        glEnableVertexAttribArray(1)
        glVertexAttribPointer(1, 3, GL_FLOAT, false, 0, 0L)

        tbo1 = glGenBuffers()
        glBindBuffer(GL_TRANSFORM_FEEDBACK_BUFFER, tbo1)
        glBufferData(GL_TRANSFORM_FEEDBACK_BUFFER, (pointCount * 28/*7*/).toLong(), GL_DYNAMIC_COPY)

        vertex1 = glGenVertexArrays()
        glBindVertexArray(vertex1)
        glEnableVertexAttribArray(0)
        glVertexAttribPointer(0, 4, GL_FLOAT, false, 0, 0L)
        glEnableVertexAttribArray(1)
        glVertexAttribPointer(1, 3, GL_FLOAT, false, 0, 0L)

        bufferGeometry = glGenBuffers()
        glBindBuffer(GL_TRANSFORM_FEEDBACK_BUFFER, bufferGeometry)
        glEnable(GL_TEXTURE_2D)
        texture = glCreateTextures(GL_TEXTURE_2D)
        glBufferData(GL_TRANSFORM_FEEDBACK_BUFFER, 16777216L /*1024*1024*4*4*/, GL_DYNAMIC_COPY)
        glTexBuffer(GL_TEXTURE_BUFFER, GL_RGB32F, bufferGeometry)

        vertexRender = glGenVertexArrays()
        glBindBuffer(GL_TRANSFORM_FEEDBACK_BUFFER, bufferGeometry)
        glEnableVertexAttribArray(0)
        glVertexAttribPointer(0, 4, GL_FLOAT, false, 0, 0L)

        val bytes = Resource.getBytes("vbm/armadillo_low.vbm")
        vbm = VBMObject(ByteBuffer.wrap(bytes))
        vbmVertex = vbm.createVAO()

        feedbacks = glGenTransformFeedbacks()
        glBindTransformFeedback(GL_TRANSFORM_FEEDBACK, feedbacks)
    }

    override fun resume() {
        super.resume()
        glClearColor(0f, 0f, 0f, 1f)
        glClearDepth(1.0)
        glEnable(GL_CULL_FACE)
        glEnable(GL_DEPTH_TEST)
        glDepthFunc(GL_LEQUAL)

        glUseProgram(programUpdate)
        uniformUpdModelMatrix = glGetUniformLocation(programUpdate, "model_matrix")
        uniformUpdProjectionMatrix = glGetUniformLocation(programUpdate, "projection_matrix")
        uniformUpdTriangleCount = glGetUniformLocation(programUpdate, "triangle_count")
        uniformUpdTimeStep = glGetUniformLocation(programUpdate, "time_step")

        glUseProgram(programRender)
        uniformRndModelMatrix = glGetUniformLocation(programRender, "model_matrix")
        uniformRndProjectionMatrix = glGetUniformLocation(programRender, "projection_matrix")
    }

    var q = 0f
    var frameCount = 0
    val x = Vector3f(1f, 0f, 0f)
    val y = Vector3f(0f, 1f, 0f)
    val z = Vector3f(0f, 0f, 1f)

    override fun update(): Boolean {
        super.update()
        val t = ((System.currentTimeMillis() and 0x3FFFF) / 0x3FFFF).toFloat()
        val projMatrix = Matrix4f().identity().frustum(-1f, 1f, -aspect, aspect, 1f, 5000f)
        val modelMatrix = Matrix4f().identity()
            .scale(0.3f)
            .rotate(t * 360f, 0f, 1f, 0f)
            .rotate(t * 360f * 3f, 0f, 0f, 1f)

        glUseProgram(programRender)
        glUniformMatrix4fv(uniformRndModelMatrix, false, modelMatrix[FloatArray(16)])
        glUniformMatrix4fv(uniformRndProjectionMatrix, false, projMatrix[FloatArray(16)])
        glBindVertexArray(vertexRender)
        glBindBuffer(GL_TRANSFORM_FEEDBACK_BUFFER, 0)
        glBeginTransformFeedback(feedbacks)
        glBindBuffer(GL_ARRAY_BUFFER, vbmVertex.buffer)
        glBindVertexArray(vbmVertex.vertex)
        glEnableVertexAttribArray(0)
        glEnableVertexAttribArray(1)
        glEnableVertexAttribArray(2)
        if (vbm.header.indexCount > 0) {
            glDrawElements(GL_TRIANGLES, vbm.frames[0].count, GL_UNSIGNED_INT, vbm.frames[0].first * 4L)
        } else {
            glDrawArrays(GL_TRIANGLES, 0, vbm.frames[0].count)
        }
        glEndTransformFeedback()

        glUseProgram(programUpdate)
        glUniformMatrix4fv(uniformUpdModelMatrix, false, Matrix4f().identity()[FloatArray(16)])
        glUniformMatrix4fv(uniformUpdProjectionMatrix, false, projMatrix[FloatArray(16)])
        glUniform1i(uniformUpdTriangleCount, vbm.frames[0].count / 3)
        if (t > q) {
            glUniform1f(uniformUpdTimeStep, (t - q) * 2000f)
        }
        q = t;
        if (frameCount and 1 != 0) {
            glBindVertexArray(vertex1)
        } else {
            glBindVertexArray(vertex0)
        }
        glBindBuffer(GL_TRANSFORM_FEEDBACK_BUFFER, 0)
        glBeginTransformFeedback(GL_POINTS)
        glDrawArrays(GL_POINTS, 0, min(pointCount, frameCount shr 3))
        glEndTransformFeedback()
        glBindVertexArray(0)
        frameCount++
        return false
    }

    override fun stop() {
        super.stop()
        // vertex
        glBindVertexArray(0)
        // buffer
        glBindBuffer(GL_TRANSFORM_FEEDBACK_BUFFER, 0)
        glBindBuffer(GL_ARRAY_BUFFER, 0)
        glBindBuffer(GL_TEXTURE_BUFFER, 0)
        // program
        glUseProgram(0)
    }

    override fun terminal() {
        // vertex
        glDeleteVertexArrays(vertexRender)
        glDeleteVertexArrays(vertex0)
        glDeleteVertexArrays(vertex1)
        glDeleteVertexArrays(vbmVertex.vertex)
        // buffer
        glDeleteBuffers(tbo0)
        glDeleteBuffers(tbo1)
        glDeleteBuffers(bufferGeometry)
        glDeleteBuffers(vbmVertex.buffer)
        // program
        glDeleteProgram(programRender)
        glDeleteProgram(programUpdate)
    }

    override fun onResize(width: Int, height: Int) {
        aspect = height / width.toFloat()
    }
}

val TransformFeedbacksUpdateVertex = """
#version 410
uniform mat4 model_matrix;
uniform mat4 projection_matrix;
uniform int triangle_count;

layout (location = 0) in vec4 position;
layout (location = 1) in vec3 velocity;

out vec4 position_out;
out vec3 velocity_out;

uniform samplerBuffer geometry_tbo;
uniform float time_step = 0.02;

bool intersect(vec3 origin, vec3 direction, vec3 v0, vec3 v1, vec3 v2, out vec3 point){
    vec3 u, v, n;
    vec3 w0, w;
    float r, a, b;

    u = (v1 - v0);
    v = (v2 - v0);
    n = cross(u, v);

    w0 = origin - v0;
    a = -dot(n, w0);
    b = dot(n, direction);

    r = a / b;
    if (r < 0.0 || r > 1.0)
        return false;

    point = origin + r * direction;

    float uu, uv, vv, wu, wv, D;

    uu = dot(u, u);
    uv = dot(u, v);
    vv = dot(v, v);
    w = point - v0;
    wu = dot(w, u);
    wv = dot(w, v);
    D = uv * uv - uu * vv;

    float s, t;

    s = (uv * wv - vv * wu) / D;
    if (s < 0.0 || s > 1.0)
        return false;
    t = (uv * wu - uu * wv) / D;
    if (t < 0.0 || (s + t) > 1.0)
        return false;

    return true;
}

vec3 reflect_vector(vec3 v, vec3 n)
{
    return v - 2.0 * dot(v, n) * n;
}

void main(void)
{
    vec3 accelleration = vec3(0.0, -0.3, 0.0);
    vec3 new_velocity = velocity + accelleration * time_step;
    vec4 new_position = position + vec4(new_velocity * time_step, 0.0);
    vec3 v0, v1, v2;
    vec3 point;
    int i;
    for (i = 0; i < triangle_count; i++)
    {
        v0 = texelFetch(geometry_tbo, i * 3).xyz;
        v1 = texelFetch(geometry_tbo, i * 3 + 1).xyz;
        v2 = texelFetch(geometry_tbo, i * 3 + 2).xyz;
        if (intersect(position.xyz, position.xyz - new_position.xyz, v0, v1, v2, point))
        {
            vec3 n = normalize(cross(v1 - v0, v2 - v0));
            new_position = vec4(point + reflect_vector(new_position.xyz - point, n), 1.0);
            new_velocity = 0.8 * reflect_vector(new_velocity, n);
        }
    }
    if (new_position.y < -40.0)
    {
        new_position = vec4(-new_position.x * 0.3, position.y + 80.0, 0.0, 1.0);
        new_velocity *= vec3(0.2, 0.1, -0.3);
    }
    velocity_out = new_velocity * 0.9999;
    position_out = new_position;
    gl_Position = projection_matrix * (model_matrix * position);
}
""".trimIndent()

val TransformFeedbacksUpdateFragment = """
#version 410

layout (location = 0) out vec4 color;

void main(void)
{
    color = vec4(1.0);
}
""".trimIndent()

val TransformFeedbacksRenderVertex = """
#version 410

uniform mat4 model_matrix;
uniform mat4 projection_matrix;

layout (location = 0) in vec4 position;
layout (location = 1) in vec3 normal;

out vec4 world_space_position;

out vec3 vs_fs_normal;

void main(void)
{
    vec4 pos = (model_matrix * (position * vec4(1.0, 1.0, 1.0, 1.0)));
    world_space_position = pos;
    vs_fs_normal = normalize((model_matrix * vec4(normal, 0.0)).xyz);
    gl_Position = projection_matrix * pos;
}
""".trimIndent()

val TransformFeedbacksRenderFragment = """
#version 410

layout (location = 0) out vec4 color;

in vec3 vs_fs_normal;

void main(void)
{
    color = vec4(0.0, 0.2, 0.0, 1.0) + vec4(0.2, 0.5, 0.4, 1.0) * abs(vs_fs_normal.z) + vec4(0.8, 0.9, 0.7, 1.0) * pow(abs(vs_fs_normal.z), 70.0);
}
""".trimIndent()